import argparse


def get_train_args() -> argparse.Namespace:
    """
    Training Argument Parser.
    Returns:
        Namespace: Parsed arguments.
    """

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--dataset_name",
        type=str,
        default="SciERC",
        help="dataset name",
    )

    parser.add_argument(
        "--training_num",
        type=int,
        default="10",
        help="从训练集中抽取的样本数",
    )

    parser.add_argument(
        "--answer_template",
        type=str,
        default="triple",
        choices=["triple", "svo", "template"],
        help="格式化答案的风格，分别是三元组，主谓宾结构和自定义the relation between",
    )

    parser.add_argument(
        "--replace_label_to_special_token",
        type=bool,
        default=False,
        help="是否使用special token替换label",
    )

    parser.add_argument(
        "--context_format",
        type=str,
        default="tag",
        choices=["tag", "conditions"],
        help="格式化句子的风格，分别是使用标签包裹实体，和将其作为条件，放在句子后面",
    )

    parser.add_argument(
        "--path_to_instructions",
        type=str,
        default="instructions.json",
        help="file with instruction prompts",
    )

    parser.add_argument(
        "--root_data_dir",
        type=str,
        default="/home/wangxiaoli/datasets/IE_INSTRUCTIONS/RE",
        help="root data directory",
    )

    parser.add_argument("--tasks", type=str, help="辅助任务列表")

    parser.add_argument(
        "--few_shot_num",
        type=int,
        default=50,
        help="few shot的数量",
    )

    parser.add_argument(
        "--log_dir",
        type=str,
        default="runs_re",
        help="where to log tensorboard",
    )

    parser.add_argument(
        "--eval_every_n_batches",
        type=int,
        default=500,
        help="do evaluation every n batches",
    )

    parser.add_argument(
        "--pred_every_n_batches",
        type=int,
        default=500,
        help="write random sample sample predictions every n batches",
    )

    parser.add_argument(
        "--path_to_model_config",
        type=str,
        required=True,
        default="configs/config_conll04.yaml",
        help="path to all necessary information for model",
    )

    parser.add_argument(
        "--path_to_model_save",
        type=str,
        default="checkpoint/",
        help="where to save model",
    )

    parser.add_argument(
        "--merge_dataset_dir",
        type=str,
        help="合并数据集的目录",
    )

    parser.add_argument(
        "--rules",
        type=str,
        help="用于增强数据集的规则",
    )

    parser.add_argument("--seed", type=int, default=42, help="random seed")

    parser.add_argument(
        "--aug_times",
        type=int,
        default=1,
        help="数据增强的倍数",
    )

    args = parser.parse_args()

    return args


def get_data_args() -> argparse.Namespace:
    """
    Reader Argument Parser.
    Returns:
        Namespace: Parsed arguments.
    """

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--path_to_file",
        type=str,
        required=True,
        help="path to initial raw data file",
    )

    parser.add_argument(
        "--dataset_type",
        type=str,
        required=True,
        choices=["conll2003", "spacy", "mit"],
        help="dataset type to map it with relevant Reader",
    )

    parser.add_argument(
        "--output_folder",
        type=str,
        help="where to save converted dataset",
    )

    args = parser.parse_args()

    return args


def get_evaluate_args() -> argparse.Namespace:
    """
    Evaluation Argument Parser.
    Returns:
        Namespace: Parsed arguments.
    """

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--model_path_or_name",
        type=str,
        required=True,
        # default="olgaduchovny/t5-base-qa-ner-conll",
        default="checkpoint",
        help="path to trained model or HF model name",
    )

    parser.add_argument(
        "--path_to_model_config",
        type=str,
        required=True,
        default="configs/config_conll04.yaml",
        help="path to all necessary information for model",
    )

    parser.add_argument(
        "--path_to_options",
        type=str,
        default="options_re.json",
        help="file with mapping dataset to its entities",
    )

    parser.add_argument(
        "--path_to_instructions",
        type=str,
        default="instructions_re.json",
        help="file with instruction prompts",
    )

    args = parser.parse_args()

    return args
